package hbase.template;

import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.*;

import java.io.IOException;
import java.nio.charset.Charset;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

public class HBaseTemplate implements HBaseOperations {
    private static final Charset charset = Charset.forName("UTF-8");
    private HBaseConnectionFactory connectionFactory;
    private ExecutorService executorService;

    public HBaseTemplate(HBaseConnectionFactory connectionFactory, ExecutorService executorService) {
        this.connectionFactory = connectionFactory;
        this.executorService = executorService;
    }

    @Override
    public <T> T execute(String tableName, TableCallback<T> callback) {
        try (Table table = connectionFactory.getTable(tableName, executorService)) {
            return callback.doInTable(table);
        } catch (Throwable th) {
            if (th instanceof Error) {
                throw ((Error) th);
            }
            if (th instanceof RuntimeException) {
                throw ((RuntimeException) th);
            }
            throw new RuntimeException(th);
        }
    }

    public Table getTable(String tableName) {
        return connectionFactory.getTable(tableName, executorService);
    }

    @Override
    public <T> List<T> find(String tableName, Scan scan, RowMapper<T> mapper) {
        return execute(tableName, table -> StreamSupport.stream(table.getScanner(scan).spliterator(), false)
                .map(mapper::mapRow).collect(Collectors.toList()));
    }

    @Override
    public <T> T get(String tableName, byte[] rowKey, RowMapper<T> mapper) {
        return get(tableName, rowKey, null, mapper);
    }

    @Override
    public <T> T get(String tableName, byte[] rowKey, String familyName, RowMapper<T> mapper) {
        return execute(tableName, table -> {
            Get get = new Get(rowKey);
            if (familyName != null) {
                get.addFamily(familyName.getBytes(charset));
            }
            Result result = table.get(get);
            return mapper.mapRow(result);
        });
    }

    @Override
    public <T> void put(String tableName, T record, RowSerializer<T> serializer) {
        execute(tableName, table -> {
            Put put = new Put(serializer.generateRowKey(record));
            serializer.serialize(record).forEach(cell -> {
                try {
                    put.add(cell);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            });
            table.put(put);
            return null;
        });
    }

    @Override
    public <T> void batchInsert(String tableName, List<T> records, RowSerializer<T> serializer) {
        List<Put> list = records.stream().map(r -> {
            Put put = new Put(serializer.generateRowKey(r));
            serializer.serialize(r).forEach(cell -> {
                try {
                    put.add(cell);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            });
            return put;
        }).collect(Collectors.toList());
        batch(tableName, 5, list);
    }

    @Override
    public long rowCount(String tableName, Scan scan) {
        return execute(tableName, table -> {
            ResultScanner rs = table.getScanner(scan);
            long count = 0;
            for (Result r : rs) {
                count += r.size();
            }
            return count;
        });
    }

    @Override
    public void delete(String tableName, byte[] rowKey) {
        delete(tableName, rowKey, null, null);
    }

    @Override
    public void delete(String tableName, byte[] rowKey, String familyName) {
        delete(tableName, rowKey, familyName, null);
    }

    @Override
    public void delete(String tableName, byte[] rowKey, String familyName, String qualifier) {
        execute(tableName, table -> {
            Delete delete = new Delete(rowKey);
            if (qualifier != null) {
                delete.addColumn(familyName.getBytes(charset), qualifier.getBytes(charset));
            } else if (familyName != null) {
                delete.addFamily(familyName.getBytes(charset));
            }
            table.delete(delete);
            return null;
        });
    }

    /**
     * 批量操作
     * @param tableName
     * @param bufferSize 单位MB
     * @param mutations
     */
    public void batch(String tableName, long bufferSize, List<? extends Mutation> mutations) {
        BufferedMutatorParams params = new BufferedMutatorParams(TableName.valueOf(tableName));
        params.writeBufferSize(bufferSize * 1024 * 1024);
        try (BufferedMutator mutator = connectionFactory.getConnection().getBufferedMutator(params)) {
            mutator.mutate(mutations);
            mutator.flush();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }
}
